################################################################################
#
# Software Detection based on
# Pseudo Stat Model
#
################################################################################

import pyflag.DB as DB
import pyflag.FlagFramework as FlagFramework
import pyflag.pyflaglog as pyflaglog
import pyflag.conf
from SoftwareUtil import *


config = pyflag.conf.ConfObject()
config.add_option("PS_TABLE", default = "ps_data")
config.add_option("PS_RESULT_TABLE", default = "ps_result")
LIMIT = 3


class PseudoStatTables(FlagFramework.EventHandler):
    def create(self, dbh, case):
        dbh.execute("""CREATE TABLE %s AS SELECT ProductCode, 0 as Count FROM %s.%s""",
            (config.PS_TABLE, config.NSRL_DB, config.NSRL_PRODUCT_COUNT))



def classif(case, productCode):
    dbh = DB.DBO(case)
    dbh.execute("""UPDATE %s SET Count = Count + 1 WHERE ProductCode = %s""",
        (config.PS_TABLE, productCode))


def getSumProductCode(sqlProductSet):
    dbh = DB.DBO(config.NSRL_DB)
    dbh.execute("""SELECT sum(NSRLCount) AS Sum FROM %s WHERE ProductCode IN %s""",
        (config.NSRL_PRODUCT_COUNT, sqlProductSet))
    row = dbh.fetch()
    sum = row['Sum']
    return sum


def calculate(case, sqlProductSet, c):
    ## Calculate a priori
    sum = getSumProductCode(sqlProductSet)
    dbh = DB.DBO(case)

    ## Important query
    dbh.execute("""SELECT
            A.ProductCode,
            (1 - (A.NSRLCount / %s)) AS a_priori,
            (1 - (A.NSRLCount / %s)) * (B.Count / A.NSRLCount) AS S
        FROM %s.%s A INNER JOIN %s B ON A.ProductCode = B.ProductCode
        WHERE A.ProductCode IN %s ORDER BY S DESC, a_priori DESC""", (sum, sum,
        config.NSRL_DB, config.NSRL_PRODUCT_COUNT, config.PS_TABLE, sqlProductSet))
    row = dbh.fetch()
    #return row['ProductCode']
    count = 0

    while row and count < c and count < LIMIT:
	yield row['ProductCode']
	row = dbh.fetch()
	count += 1




def pseudo_stat_calculate(case):
    pyflaglog.log(pyflaglog.DEBUG, "Sorting hashes...")
    dbh = DB.DBO(case)
    dbh.execute("""SELECT SHA1, MD5, Count FROM %s WHERE N > 0
        ORDER BY N""", config.SOFT_SCANNER_HASHES)
    row = dbh.fetch()

    ## Process all hashes
    pyflaglog.log(pyflaglog.DEBUG, "Processing hashes...")

    while row:
        sha1 = row['SHA1']; md5 = row['MD5']; c = int(row['Count']);
        productCodeList = getProductCodeList(sha1, md5)
        sqlProductSet = getSQLProductSet(productCodeList)
	for productCode in calculate(case, sqlProductSet, c):
    	    classif(case, productCode)

	## Get next hash
        row = dbh.fetch()

    ## Creating result table
    pyflaglog.log(pyflaglog.INFO, "Creating PS result table...")
    dbh.execute("""CREATE TABLE %s AS SELECT DISTINCT
            A.ProductCode,
	    C.ProductName,
	    C.ProductVersion,
            A.Count,
            B.NSRLCount,
            round(A.Count / B.NSRLCount, 2) AS Result
        FROM %s A
            INNER JOIN %s.%s B ON A.ProductCode = B.ProductCode
	    INNER JOIN %s.%s C ON B.ProductCode = C.ProductCode
        WHERE A.Count > 0 ORDER BY Result DESC""", (
            config.PS_RESULT_TABLE, config.PS_TABLE,
            config.NSRL_DB, config.NSRL_PRODUCT_COUNT,
	    config.NSRL_DB, config.NSRL_PRODUCTS))

